import torch
from torch.utils.data import DataLoader, TensorDataset

def prepare_data(num_samples, input_dim, num_classes):
    """
    模拟数据准备函数。
    
    Args:
    - num_samples: 数据集中的样本数。
    - input_dim: 输入特征的维度。
    - num_classes: 类别数。
    
    Returns:
    - train_loader: 训练数据加载器。
    """
    # 模拟随机生成一些数据
    # 注意：在实际应用中，这里应该是从真实数据集中加载数据
    inputs = torch.randn(num_samples, input_dim)  # 随机生成输入数据
    targets = torch.randint(0, num_classes, (num_samples,))  # 随机生成目标标签

    # 将数据和标签封装成TensorDataset
    dataset = TensorDataset(inputs, targets)

    # 创建DataLoader来批处理数据
    # 在分布式训练中，batch_size和num_workers可能需要根据GPU数量和内存进行调整
    train_loader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=0)

    return train_loader

# 在主程序中调用数据准备函数
if __name__ == "__main__":
    num_samples = 100  # 假设有100个样本
    input_dim = 10    # 假设输入特征维度为10
    num_classes = 2   # 假设是一个二分类问题

    train_loader = prepare_data(num_samples, input_dim, num_classes)

    # 接下来是模型定义、训练循环等步骤...
    # 但由于这里只关注数据准备部分，所以省略了这些步骤
